import json
import time
# 多线程
from threading import Thread

import cv2
import gym
import numpy as np
import win32con
import win32gui
import win32ui
from pykeyboard import PyKeyboard
from redis import Redis


# from pymouse import PyMouse


class HKPoolEnv2D(gym.Env):
    def __init__(self):
        self.r = Redis(host='127.0.0.1', port=6379, db=1)
        self.hwnd_id = []
        self.k = PyKeyboard()
        self.initailize()
        win32gui.EnumWindows(self.get_all_hwnd, 0)

        # 激活窗口
        win32gui.SetForegroundWindow(self.hwnd_id[0])
        # 移动窗口到左上角
        win32gui.MoveWindow(self.hwnd_id[0], -10, -31, 1280, 720, True)
        # 创建dc, 截取桌面
        self.hdesktop = win32gui.GetDesktopWindow()
        # 获取窗口大小
        l, t, r, b = win32gui.GetWindowRect(self.hwnd_id[0])
        self.width = r - l
        self.height = b - t

        # 创建设备描述表
        self.desktop_dc = win32gui.GetWindowDC(self.hdesktop)
        # 创建dc对象
        self.mfc_dc = win32ui.CreateDCFromHandle(self.desktop_dc)
        # 创建内存设备描述表
        self.save_dc = self.mfc_dc.CreateCompatibleDC()
        # 创建位图对象
        self.save_bit_map = win32ui.CreateBitmap()
        # 截图
        self.save_bit_map.CreateCompatibleBitmap(self.mfc_dc, self.width, self.height)

    def initailize(self):
        pass

    def get_all_hwnd(self, hwnd, mouse):
        if (win32gui.IsWindow(hwnd) and
            win32gui.IsWindowEnabled(hwnd) and
            win32gui.IsWindowVisible(hwnd)) and win32gui.GetWindowText(hwnd).__contains__('Hollow Knight'):
            self.hwnd_id.append(hwnd)

    def restart_game(self):
        win32gui.EnumWindows(self.get_all_hwnd, 0)
        time.sleep(1)
        # 激活窗口
        win32gui.SetForegroundWindow(self.hwnd_id[0])
        time.sleep(5)
        self.k.release_key(self.k.left_key)
        self.k.release_key(self.k.right_key)
        self.k.release_key(self.k.up_key)
        self.k.release_key(self.k.down_key)
        self.k.release_key("Z")
        self.k.release_key("X")
        self.k.release_key("C")
        self.k.release_key("F")
        time.sleep(10)
        self.k.tap_key(self.k.left_key)
        time.sleep(4)
        self.k.tap_key(self.k.up_key)
        time.sleep(2)
        self.k.tap_key("z")

    def reset(self, *args, **kwargs):
        self.no = 0
        self.enemy_hp = 100000
        self.knight_hp = 100000
        self.last_mp = 0
        self.max_x = 0
        self.min_x = 1000
        self.max_y = 0
        self.min_y = 1000
        t = Thread(target=self.restart_game)
        t.start()
        self.r.set("can_action", "1")
        state_json = self.r.brpop(['hollow_obs'])[1]
        state = json.loads(state_json)
        self.state_to_obs(state_json, True)
        binary = self.get_binary()
        print("reset")

        # 重复20帧图像 (720, 1280) -> (720, 1280, 20)
        # self.binary_20 = np.repeat(binary[:, :, np.newaxis], 20, axis=2)
        # # 每隔5帧取一帧图像 (720, 1280, 20) -> (720, 1280, 4)
        # binary_5 = self.binary_20[:, :, ::5]
        # print(binary_5.shape)
        return binary

    def act(self, move, attack, do_attack=True):
        if do_attack:
            self.do_attack(attack)
        self.do_move(move)

    def do_move(self, move):
        L = self.k.left_key
        R = self.k.right_key
        J = "Z"
        # 0: stay, 1: jump, 2: left, 3: right, 4: left jump, 5 right jump, 6 double jump
        if move == 0:
            # 释放左右和跳跃键
            self.k.release_key(L)
            self.k.release_key(R)
            self.k.release_key(J)
        if move == 1:
            self.k.release_key(L)
            self.k.release_key(R)
            self.k.press_key(J)
        if move == 2:
            self.k.release_key(R)
            self.k.release_key(J)
            self.k.press_key(L)
        if move == 3:
            self.k.release_key(L)
            self.k.release_key(J)
            self.k.press_key(R)
        if move == 4:
            self.k.release_key(R)
            self.k.press_key(L)
            self.k.press_key(J)
        if move == 5:
            self.k.release_key(L)
            self.k.press_key(R)
            self.k.press_key(J)
        if move == 6:
            self.k.release_key(J)
            self.k.press_key(J)

    def do_attack(self, attack):
        L = self.k.left_key
        R = self.k.right_key
        U = self.k.up_key
        D = self.k.down_key
        S = "F"
        C = "C"
        A = "X"
        # 0: 无动作, 1:左攻击, 2:右攻击, 3: 上攻击, 4: 下攻击, 5: 下砸, 6: 吼, 7: 左波, 8: 右波, 9: 左冲, 10: 右冲
        if attack in [1, 7, 9]:
            self.k.release_key(R)
            self.k.release_key(U)
            self.k.release_key(D)
            self.k.press_key(L)
        if attack in [2, 8, 10]:
            self.k.release_key(U)
            self.k.release_key(D)
            self.k.release_key(L)
            self.k.press_key(R)
        if attack in [3, 6]:
            self.k.release_key(D)
            self.k.press_key(U)
        if attack in [4, 5]:
            self.k.release_key(U)
            self.k.press_key(D)
        if attack in [1, 2, 3, 4]:
            self.k.tap_key(A)
        if attack in [5, 6, 7, 8]:
            self.k.tap_key(S)
        if attack in [9, 10]:
            self.k.tap_key(C)
        self.k.release_key(L)
        self.k.release_key(R)
        self.k.release_key(U)
        self.k.release_key(D)

    def get_binary(self, save=False, no=0):
        win32gui.EnumWindows(self.get_all_hwnd, 0)

        # 激活窗口
        win32gui.SetForegroundWindow(self.hwnd_id[0])
        # 移动窗口到左上角
        win32gui.MoveWindow(self.hwnd_id[0], -10, -31, 1280, 720, True)

        self.save_dc.SelectObject(self.save_bit_map)
        # 截取从左上角(50, 50)长宽为(width, height)的图片
        self.save_dc.BitBlt((0, 0), (self.width - 21, self.height - 41), self.mfc_dc, (0, 0), win32con.SRCCOPY)
        # 保存位图文件
        signedIntsArray = self.save_bit_map.GetBitmapBits(True)
        img = np.frombuffer(signedIntsArray, dtype='uint8')
        img.shape = (self.height, self.width, 4)
        # 缩放
        img = cv2.resize(img, (255, 255))
        if save:
            cv2.imwrite("imgs/{}.png".format(no), img)
        # 保存
        # cv2.imwrite('test.png', img)
        # 转换为灰度图
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        if save:
            cv2.imwrite("imgs/gray_{}.png".format(no), gray)
        # 压缩到0-1
        gray = gray / 255.0
        # 缩放
        # gray = cv2.resize(gray, (511, 511))
        # 保存
        # cv2.imwrite('test_gray.png', gray)
        # 二值化
        # ret, binary = cv2.threshold(gray, 127, 1, cv2.THRESH_BINARY)
        # if save:
        #     ret, binary_255 = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
        #     cv2.imwrite("imgs/binary_{}.png".format(no), binary_255)
        # 转成numpy数组
        binary = np.array(gray)
        return binary

    def step(self, *args, **kwargs):
        reward_m = -abs(self.enemy_x - self.knight_x) / 10
        reward_a = -0.01
        self.act(args[0], args[1])
        if reward_a == 0:
            reward_a -= 0.1
        # 如果boss在右边
        if abs(self.enemy_x - self.knight_x) > 2:
            if self.enemy_x > self.knight_x:
                # 向右走
                if args[0] in [3, 5]:
                    reward_m += 1 + abs(self.enemy_x - self.knight_x) / 100
                # 向左走
                if args[0] in [2, 4]:
                    reward_m -= 1
                # 左攻击
                # if args[1] == 1:
                #     reward_a -= 1
            # 如果boss在左边,且
            if self.enemy_x < self.knight_x:
                # 向左走
                if args[0] in [2, 4]:
                    reward_m += 1 + abs(self.enemy_x - self.knight_x) / 100
                # 向右走
                if args[0] in [3, 5]:
                    reward_m -= 5
                # 右攻击
                # if args[1] == 2:
                #     reward_a -= 1
        # if self.distance < 2:
        #     if args[1] in [9, 10]:
        #         reward_a += 1
        if self.knight_mp < 33:
            if args[1] in [5, 6, 7, 8]:
                reward_a -= 1
        while True:
            binary = self.get_binary()
            self.r.set("can_action", "1")
            state_json = self.r.brpop(['hollow_obs'])[1]
            self.state_to_obs(state_json)
            if self.done:
                break
            if self.knight_damage > 0 or self.enemy_damage > 0:
                break
            if not self.knight_spell or not self.dashing:
                break
        # binary = self.get_binary()
        # self.r.set("can_action", "1")
        # state_json = self.r.brpop(['hollow_obs'])[1]
        # self.state_to_obs(state_json)
        if 3 < self.distance < 4:
            reward_m += 1
        if self.enemy_damage > 0:
            reward_m += self.enemy_damage / 10
            reward_a += self.enemy_damage / 4
            # if args[1] in [5, 6, 7, 8]:
            #     reward_a += self.enemy_damage / 5
        if self.knight_damage > 0:
            if args[1] in [5, 6, 7, 8, 9, 10]:
                reward_a = -2
            else:
                reward_m = -2
                reward_a = -1
        if self.done:
            reward_a += 1
            reward_m += 1
            self.k.release_key(self.k.left_key)
            self.k.release_key(self.k.right_key)
            self.k.release_key(self.k.up_key)
            self.k.release_key(self.k.down_key)
            self.k.release_key("Z")
            self.k.release_key("X")
            self.k.release_key("C")
            self.k.release_key("F")
        print(f'\r##max_x:{self.max_x}, min_x:{self.min_x}, max_y:{self.max_y}, min_y:{self.min_y}##', end=" ")
        return binary, reward_m, reward_a, self.done, {}

    def state_to_obs(self, state, flag=False):
        state_dict = json.loads(state)
        # 是否黑冲
        self.shadow_dashing = True if state_dict['shadow_dashing'] == '1' else False
        # 黑冲时间
        self.shadow_dash_time = float(state_dict['shadow_dash_time'])

        # 是否冲刺
        self.dashing = True if state_dict['dashing'] == '1' else False
        # 是否在攻击
        self.is_attacking = True if state_dict['is_attacking'] == '1' else False
        # 是否在释放技能
        self.knight_spell = True if state_dict['knight_spell'] != "Inactive" else False

        # boss 和 knight 的伤害
        self.knight_damage = float(state_dict['knight_damage'])
        self.enemy_damage = float(state_dict['enemy_damage'])
        # 游戏是否结束
        self.done = int(state_dict['done'])

        # knight 的位置 x 减掉 70 y 减掉 10
        self.knight_x = float(state_dict['knight_x'])
        self.knight_y = float(state_dict['knight_y'])

        # knight 的灵魂量
        self.knight_mp = float(state_dict['knight_mp'])

        # boss 的位置 x 减掉 70 y 减掉 10
        self.enemy_x = float(state_dict['enemy_x'])
        self.enemy_y = float(state_dict['enemy_y'])

        # boss和knight的向量
        self.enemy_knight_x = self.enemy_x - self.knight_x
        self.enemy_knight_y = self.enemy_y - self.knight_y

        # 距离
        self.distance = np.sqrt(self.enemy_knight_x ** 2 + self.enemy_knight_y ** 2)
        # self.distance = np.array([self.distance], dtype=np.float32)
        # 角度,0-360
        self.angle = np.arctan2(self.enemy_knight_y, self.enemy_knight_x) * 180 / np.pi
        self.angle = self.angle + 360 if self.angle < 0 else self.angle

        # # self.angle = np.arctan2(self.enemy_knight_y, self.enemy_knight_x)
        if flag:
            # 重复20帧
            self.knight_x_20 = np.tile(np.array([self.knight_x]), 20)
            self.knight_y_20 = np.tile(np.array([self.knight_y]), 20)
            self.enemy_x_20 = np.tile(np.array([self.enemy_x]), 20)
            self.enemy_y_20 = np.tile(np.array([self.enemy_y]), 20)
            self.enemy_knight_x_20 = np.tile(np.array([self.enemy_knight_x]), 20)
            self.enemy_knight_y_20 = np.tile(np.array([self.enemy_knight_y]), 20)
        else:
            # 去掉最旧帧,添加最新帧
            self.knight_x_20 = np.concatenate((self.knight_x_20[1:], np.array([self.knight_x])))
            self.knight_y_20 = np.concatenate((self.knight_y_20[1:], np.array([self.knight_y])))
            self.enemy_x_20 = np.concatenate((self.enemy_x_20[1:], np.array([self.enemy_x])))
            self.enemy_y_20 = np.concatenate((self.enemy_y_20[1:], np.array([self.enemy_y])))
            self.enemy_knight_x_20 = np.concatenate((self.enemy_knight_x_20[1:], np.array([self.enemy_knight_x])))
            self.enemy_knight_y_20 = np.concatenate((self.enemy_knight_y_20[1:], np.array([self.enemy_knight_y])))
        # 每隔5帧取一帧
        self.knight_x_4 = self.knight_x_20[::5]
        self.knight_y_4 = self.knight_y_20[::5]
        self.enemy_x_4 = self.enemy_x_20[::5]
        self.enemy_y_4 = self.enemy_y_20[::5]
        # 速度 第一帧减去最后一帧
        self.enemy_speed_x = self.enemy_x_4[-1] - self.enemy_x_4[0]
        self.enemy_speed_y = self.enemy_y_4[-1] - self.enemy_y_4[0]
        self.enemy_knight_x_4 = self.enemy_knight_x_20[::5]
        self.enemy_knight_y_4 = self.enemy_knight_y_20[::5]
        # boss x速度是否指向knight, float32
        self.enemy_speed_x_knight = True if self.enemy_speed_x * self.enemy_knight_x_4[0] > 0 else False
        # boss y速度是否指向knight
        self.enemy_speed_y_knight = True if self.enemy_speed_y * self.enemy_knight_y_4[0] > 0 else False

        # 计算boss速度和knight的向量的夹角
        if self.enemy_speed_x == 0:
            self.enemy_speed_x = 0.000001
        if self.enemy_speed_y == 0:
            self.enemy_speed_y = 0.000001
        if self.enemy_knight_x == 0:
            self.enemy_knight_x = 0.000001
        if self.enemy_knight_y == 0:
            self.enemy_knight_y = 0.000001
        self.enemy_speed = [self.enemy_speed_x, self.enemy_speed_y]
        self.enemy_knight = [-self.enemy_knight_x, -self.enemy_knight_y]
        self.enemy_speed_knight_angle = np.arccos(
            np.dot(self.enemy_speed, self.enemy_knight) / (
                    np.linalg.norm(self.enemy_speed) * np.linalg.norm(self.enemy_knight))) * 180 / np.pi
        # # 将所有状态拼接起来, extend
        # obs = np.concatenate((self.onGround, self.isJumping, self.isDoubleJumping, self.isFalling,
        #                       self.isInvincible, self.shadow_dashing, self.dashing, self.is_attacking,
        #                       self.knight_spell,
        #                       self.knight_x_4, self.knight_y_4, self.knight_mp, self.enemy_x_4, self.enemy_y_4,
        #                       self.enemy_knight_x_4, self.enemy_knight_y_4, self.enemy_speed_y_knight,
        #                       self.enemy_speed_x_knight,
        #                       self.shadow_dash_time))
        if self.enemy_x < self.min_x:
            self.min_x = self.enemy_x
        if self.knight_x < self.min_x:
            self.min_x = self.knight_x
        if self.enemy_x > self.max_x:
            self.max_x = self.enemy_x
        if self.knight_x > self.max_x:
            self.max_x = self.knight_x
        if self.enemy_y > self.max_y:
            self.max_y = self.enemy_y
        if self.knight_y > self.max_y:
            self.max_y = self.knight_y
        if self.enemy_y < self.min_y:
            self.min_y = self.enemy_y
        if self.knight_y < self.min_y:
            self.min_y = self.knight_y
        return


if __name__ == '__main__':
    env = HKPoolEnv2D()
    # env.restart_game()z
    win32gui.EnumWindows(env.get_all_hwnd, 0)
    handle = env.hwnd_id[0]
    # 激活窗口
    win32gui.SetForegroundWindow(handle)
    time.sleep(3)
    k = PyKeyboard()
    k.press_key("Z")
    time.sleep(0.2)
    k.release_key("Z")
    time.sleep(0.1)
    k.press_key("Z")
    time.sleep(1)
    k.release_key('Z')

    obs = env.reset()
    print(obs)
    for i in range(100):
        action = np.random.randint(0, 12)
        obs, reward, done, info = env.step(1, 1)
        count = 0
        print(obs, reward, done, info)
